import gradio as gr
from sklearn.cluster import KMeans
import numpy as np
from PIL import Image
import tempfile
import os

def extract_colors(image, num_colors=5):
    # 将numpy数组转换为PIL图像
    img = Image.fromarray(image)
    img = img.resize((64, 64))  # 缩小图片以减少计算量
    img_array = np.array(img)
    # 将图片转换为一维数组，每个像素点的RGB值
    pixels = img_array.reshape(-1, 3)
    # 使用KMeans算法提取颜色
    kmeans = KMeans(n_clusters=num_colors, random_state=0).fit(pixels)
    colors = kmeans.cluster_centers_
    
    # 将颜色转换为十六进制
    hex_colors = ["#"+"".join([format(int(i), '02x') for i in color]) for color in colors]
    
    # 创建颜色的图像并保存为临时文件
    color_images = []
    for i, color in enumerate(colors):
        color_img = Image.new('RGB', (100, 100), color=tuple(int(c) for c in color))
        with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as tmp:
            color_img.save(tmp.name)
            color_images.append(tmp.name)
    
    return color_images, hex_colors

# 创建Gradio界面
iface = gr.Interface(
    fn=extract_colors,
    inputs=gr.Image(label="上传图片"),
    outputs=[gr.Gallery(label="主色调调色板"), gr.Textbox(label="颜色的十六进制表示")],
    live=True
)

if __name__ == "__main__":
    iface.launch()